##########################################################################################

library(ComplexHeatmap)
library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)
library(circlize)

##########################################################################################

option_list <- list(
    make_option(c("--maf_path"), type = "character") ,
    make_option(c("--images_path"), type = "character") ,
    make_option(c("--info_file"), type = "character") ,
    make_option(c("--info_public_file"), type = "character") ,
    make_option(c("--class_order_file"), type = "character") 
)

if(1!=1){
    
    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic"
    maf_path <- paste(work_dir,"/","maf_public",sep="")
	images_path <- paste(work_dir,"/","images",sep="")
	info_file <- paste(work_dir,"/config/tumor_normal.class.list",sep="")
	info_public_file <- paste(work_dir,"/public_ref/combine/MutationInfo.combine.addMolecularSubType.tsv",sep="")
	class_order_file <- paste(work_dir,"/config/Class_order.list",sep="")
}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

maf_path <- opt$maf_path
info_file <- opt$info_file
info_public_file <- opt$info_public_file
images_path <- opt$images_path
class_order_file <- opt$class_order_file

###########################################################################################

im_maf_file <- paste0(maf_path , "/All_use.IM.maf")

###########################################################################################

class_order <- data.frame(fread(class_order_file , header = T))
info <- data.frame(fread(info_public_file))

dat_im <- data.frame(fread(im_maf_file))

###########################################################################################
## 去除MSI样本
info <- subset( info , !(Molecular.subtype %in% c("POLE" , "MSI") ))
info$Molecular.subtype <- ifelse( info$Molecular.subtype %in% c("EBV" , "unknown") , "Other" , info$Molecular.subtype )
## 去除IGC和DGC多个的样本
info <- subset( info , Class != "IGC + DGC" )

###########################################################################################
## 一个人多个样本算一个样本
## 只看IM
dat_im$Tumor_Sample_Barcode <- paste0( dat_im$Tumor_Sample_Barcode , "_IM" )

info_im <- subset( info , From == "NJMU"  )
info_im$Tumor <- paste0( info_im$Tumor , "_IM" )
info_im$Class <- "IM"
info_im$Stage <- "unknown"
info <- info_im
info$Normal <- info$Tumor
rownames(info) <- info$Tumor

info$Gender[info$Gender=="male"] <- "Male"
info$Gender[info$Gender=="female"] <- "Female"

info$HP <- ifelse( info$HP == "unknown" , "Unknown" , info$HP )
info$HP <- factor(info$HP , levels = c("Positive", "Negative" , "Unknown" ) , ordered=T)

info$Alcohol <- ifelse( info$Alcohol == "unknown" , "Unknown" , info$Alcohol )
info$Alcohol <- ifelse( info$Alcohol == "No" , "Non-drinker" , info$Alcohol )
info$Alcohol <- ifelse( info$Alcohol == "Drink" , "Drinker" , info$Alcohol )
info$Alcohol <- factor( info$Alcohol , levels = c("Drinker", "Non-drinker" , "Unknown" ) , ordered=T)

###########################################################################################

Variant_Types <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")
Variant_Type_Combine <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift","In_Frame","Splice_Site","Nonstop_Mutation","Multiple_Hits")


###########################################################################################

show_gene <- c("MUC6" , "CFTR" , "BMP6")
dat_im_high <- subset( dat_im , Hugo_Symbol %in% show_gene)
dat_im_high$Hugo_Symbol <- "MUC6_CFTR_BMP6"
smg <- c("MUC6_CFTR_BMP6" , show_gene)

###########################################################################################

CreateMutMatrix <- function(dat = dat , smg = smg , Variant_Type = Variant_Type ){

	mut <- dat
	mut$Variant_Classification <- as.character(mut$Variant_Classification)
	mut <- mut[which(mut$Variant_Classification %in% Variant_Type),]

	## INDEL 和 INS 合并
	mut[grep("In_Frame",mut$Variant_Classification),'Variant_Classification'] = "In_Frame"
	mut[grep("Frame_Shift",mut$Variant_Classification),'Variant_Classification'] = "Frame_Shift"

	## Gene只考虑有突变的SMG和CGC
	Gene <- smg

	## 构建矩阵
	Sample <- info[,"Tumor"]
	maf_matrix <- matrix("" , ncol = length(Sample) , nrow = length(Gene) , dimnames = list(Gene,Sample))

	for(gene in rownames(maf_matrix)){
		print(gene)

		## 突变类型
		for(tumor in colnames(maf_matrix)){
			index <- which(mut$Hugo_Symbol==gene & mut$Tumor==tumor)
			if(length(index)==0){
				var=""
			}else if(length(index)==1){
				var <- mut[index,'Variant_Classification']
			}else if(length(index)>1){
				var <- "Multiple_Hits"
			}
			print(var)
			maf_matrix[ which(rownames(maf_matrix)==gene) , which(colnames(maf_matrix)==tumor) ] <-  var
		}
	}

	return(maf_matrix)
}

MutMatrixOrder <- function( mut = mut , info = info){


	## 将突变的百分比放在基因名上
	## 百分比为人的
	mut_per <- paste0( 100 * round(
			apply(mut , 1 , function(x){length(unique(info[info$Tumor %in% names(which(x!="")) ,"Normal"]))} ) / 
			length(unique(info$Normal)) ,
		2)) 
	mut_per <- as.numeric(mut_per)

	IM_sample <- info[info$Class=="IM","Tumor"]

	mut_pre <- mut[,colnames(mut) %in% c(IM_sample)]

	NumMut_Pre <- apply(mut_pre , 1 ,function(x){
		length(unique(info[info$Tumor %in% names(which(x!="")) ,"Normal"]))}
	)
	rownames(mut) <- paste0(rownames(mut) , " ( " , mut_per , "%" , " ) ")

	MutNumMatrix <- data.frame(NumMut_Pre)

	## 基因排序
	mut <- mut[order(MutNumMatrix$NumMut_Pre , decreasing = T ),]

	## 样本排序
	## 保证前面的mut的样本顺序和info的一致
	mut <- rbind( mut , Class = info$Class )
	mut <- mut[ , order(mut["Class",] , 
		mut[1,] , mut[2,] , mut[3,] , mut[4,] ,
		colnames(mut) , decreasing = T )]

	## 去除Class
	mut <- mut[rownames(mut)!=c("Class"),]

	return(mut)
}

plotMutWaterFull <- function( MutMatrix = MutMatrix , Variant_Type_Combine = Variant_Type_Combine , images_name = images_name ){

	################################################################################################
	## 注释的名字

	## 展示出现频率 >= 1个样本的突变
	mut <- MutMatrix[apply(MutMatrix,1,function(x){length(which(x!=""))>1}),]

	## 突变矩阵排序
	mut <- MutMatrixOrder( mut = mut , info = info )
	
	## 顺序,很重要！
	sample_order <- colnames(mut)

	molecular_order <- info[sample_order,"Molecular.subtype"]
	hp_order <- info[sample_order,"HP"]
	alcohol_order <- info[sample_order,"Alcohol"]

	################################################################################################
	## 突变的颜色
	col = c(rgb(red=48,green=115,blue=186,alpha=255,max=255),
		rgb(red=236,green=27,blue=35,alpha=255,max=255),
		rgb(red=236,green=179,blue=33,alpha=255,max=255),
		rgb(red=235,green=230,blue=26,alpha=255,max=255),
		rgb(red=150,green=131,blue=189,alpha=255,max=255),
		rgb(red=65,green=174,blue=119,alpha=255,max=255))

	names(col) = c(
	  'Missense_Mutation',
	  'Nonsense_Mutation',
	  'Frame_Shift',
	  'In_Frame',
	  'Splice_Site',
	  'Multiple_Hits'
	)

	## 设置不同背景的颜色
	alter_fun = list(
	    background = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"), 
	            gp = gpar(fill = "#F2F2F2", col = NA))
	    },
	    Missense_Mutation = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"), 
	            gp = gpar(fill = col["Missense_Mutation"], col = NA))
	    },
	    Nonsense_Mutation = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"), 
	            gp = gpar(fill = col["Nonsense_Mutation"], col = NA))
	    },
	  
	    Nonstop_Mutation = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"),
	            gp = gpar(fill = col["Nonstop_Mutation"], col = NA))
	    },
	    Frame_Shift = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"),
	            gp = gpar(fill = col["Frame_Shift"], col = NA))
	    },
	    In_Frame = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"),
	            gp = gpar(fill = col["In_Frame"], col = NA))
	    },
	    Splice_Site = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"),
	            gp = gpar(fill = col["Splice_Site"], col = NA))
	    },
	    Multiple_Hits = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"),
	            gp = gpar(fill = col["Multiple_Hits"], col = NA))
	    },
	    show_legend = FALSE 
	)

	################################################################################################
	## 顶部注释 突变数量

	col_from <- c(
		rgb(red=244, green=241,blue=222,alpha=255,max=255) ,
		rgb(red=223, green=122,blue=94,alpha=255,max=255) ,
		rgb(red=60, green=64,blue=91,alpha=255,max=255) ,
		rgb(red=130, green=178,blue=154,alpha=255,max=255) ,
		rgb(red=242, green=204,blue=142,alpha=255,max=255) 
		)
	names(col_from) <- c("NJMU" , "OncoSG" ,  "TCGA" , "TMUCIH" , "Utokyo")

	col_molecular <- c(col_from[c(2,3)] , "grey")
	names(col_molecular) <- c("CIN" , "GS" , "Other")

	col_hp <- c(col_from[c(2,3)] , "grey")
	names(col_hp) <- c("Positive" , "Negative" , "Unknown")

	col_drink <- c(col_from[c(2,3)] , "grey")
	names(col_drink) <- c("Drinker" , "Non-drinker" , "Unknown")

	## 顶部注释 class
	top_annotation <- HeatmapAnnotation(
		foo = anno_empty(height = unit(2, "cm") , border =F) ,  ## 留给注释空间
		Molecular_Subtype = molecular_order ,
		HP = hp_order ,
		Drinking_status = alcohol_order ,
		col = list( 
			Molecular_Subtype = col_molecular ,
			HP = col_hp , 
			Drinking_status = col_drink
		) ,
		annotation_name_side = "left" , 
	  	border = T ,
	  	gap = unit(1, "mm") ,
	  	#show_annotation_name = c(Mut_Num = FALSE) , 
	  	annotation_name_gp = gpar(fontsize = 12),
	  	show_legend = FALSE  ## 所有的legend均后期定义
	)

	################################################################################################
	## 给沉默突变和非沉默突变加Legend

	ldg_mol = Legend(labels = names(col_molecular) , title = "Molecular_Subtype" ,
	 legend_gp = gpar(fill = col_molecular , bar_width = 1 , fontsize = 12) , ncol = 2 , 
	 gap = unit(1, "cm") ## 图例的间隔
	)

	ldg_hp = Legend(labels = names(col_hp) , title = "HP" ,
	 legend_gp = gpar(fill = col_hp , bar_width = 1 , fontsize = 12) , ncol = 2 , 
	 gap = unit(1, "cm") ## 图例的间隔
	)

	ldg_drink = Legend(labels = names(col_drink) , title = "Drinking Status" ,
	 legend_gp = gpar(fill = col_drink , bar_width = 1 , fontsize = 12) , ncol = 2 , 
	 gap = unit(1, "cm") ## 图例的间隔
	)

	## 突变类型注释
	col_Mut <- col 

	ldg_Variant = Legend(labels = names(col_Mut) , title = "Mutations" , border = "black" , ## 注释边框加黑
	 	legend_gp = gpar(fill = col_Mut , bar_width = 1 ,fontsize = 12) , ncol = 3 , 
	 	gap = unit(1, "cm") ## 图例的间隔
	)	

	lgd_all <- packLegend(ldg_Variant , ldg_mol , ldg_hp , ldg_drink ,
		column_gap = unit(1, "cm") , row_gap = unit(10, "mm") , direction = "horizontal"  )


	################################################################################################
	p <- oncoPrint(mut, name = "cases", ## 后面加分割线用的name
	    alter_fun = alter_fun, col = col, 
	    top_annotation = top_annotation ,
	    row_names_side = "left", row_names_gp = gpar(fontsize = 12) ,  ## 基因名移到左边
	    #left_annotation = left_annotation , 
	    right_annotation = rowAnnotation(foo = anno_empty(border = FALSE)), ## 不展示突变构成比
	    left_annotation = rowAnnotation(foo = anno_empty(border = FALSE)),
	    show_pct = FALSE , ##不展示百分比
	    border = TRUE,
	    row_order = 1:nrow(mut) ,
	    show_heatmap_legend = FALSE ,
	    row_title_rot = 0 , row_title_gp = gpar(fontsize = 10) , 
	    row_gap = unit(0, "points")
	)

	##
	pdf(images_name , width = 15 , height = 3)
	draw(p )
	#draw(ldg_Variant, x = unit(0.3, "npc"), y = unit(0.005, "npc"), just = c("left", "bottom"))
	draw(lgd_all, x = unit(0.2, "npc"), y = unit(0.99, "npc"), just = c("left", "top"))
	# draw(ldg_Variant, x = unit(0.99, "npc"), y = unit(0.99, "npc"), just = c("right", "top"))
	dev.off()

}


###########################################################################################
## 合并所有的MAF文件
dat <- rbind(dat_im , dat_im_high)
dat <- dat[,colnames(dat) != "Variant_Type" ]
dat$Tumor <- dat$Tumor_Sample_Barcode

## 产生输入矩阵
Variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")
Variant_Type_Combine <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift","In_Frame","Splice_Site","Nonstop_Mutation","Multiple_Hits")

MutMatrix <- CreateMutMatrix( dat = dat , smg = smg , Variant_Type = Variant_Type )

## 瀑布图
images_name <- paste0(images_path , "/Mut_WaterFall.new.IM.pdf")
plotMutWaterFull( MutMatrix = MutMatrix , Variant_Type_Combine = Variant_Type_Combine , images_name = images_name )


